/*
 * contrib/spi/autoinc.c
 */
#include "postgres.h"

#include "access/htup_details.h"
#include "catalog/pg_type.h"
#include "commands/sequence.h"
#include "commands/trigger.h"
#include "executor/spi.h"
#include "utils/builtins.h"
#include "utils/rel.h"
#include "utils/memutils.h"
#include "access/xact.h"

PG_MODULE_MAGIC;

PG_FUNCTION_INFO_V1(autoinc);
PG_FUNCTION_INFO_V1(test_spt_feature);

char	*sql0 = "DROP TABLE IF EXISTS t1";
char	*sql1 = "CREATE TABLE t1(i int, j int, k varchar)";
char	*sql2 = "SELECT * FROM t1";

static void spi_simple_test(void);
static void outExecResult(int ret);
static void tx_begin(MemoryContext	oldcontext);
static void tx_commit(MemoryContext	oldcontext, ResourceOwner oldowner);
static void tx_abort(MemoryContext	oldcontext, ResourceOwner oldowner);
static void spi_with_args_test(void);
static void my_exec_spi_with_try(char *sql);

static void output_str(const char *fmt,...) __attribute__((format(gnu_printf,1,2)));


static void
output_str(const char *fmt,...)
{
	va_list		args;
	FILE		*temp = stdout;

	va_start(args, fmt);
	vfprintf(temp, _(fmt), args);
	va_end(args);
	fputs("\n", temp);

	fflush(temp);
}

/*
 * This function is for we can output a exec result
 * of SPI_exec()
 */
static void
outExecResult(int ret)
{
	int64		rows = 0;
	int			i = 0;
	int		j = 0;

	if(ret <= 0)
		return;
	rows = SPI_processed;
	if (SPI_tuptable)
    {
        TupleDesc 		tupdesc = SPI_tuptable->tupdesc;
        SPITupleTable	*tuptable = SPI_tuptable;
        char			buf[8192] = {0};
        

		for (i = 1; i <= tupdesc->natts; i++)
		{
			char	*value = NULL;
			char	*type = NULL;
			char	*inter = (i == tupdesc->natts) ? " " : " |";
			char	*cur_loc = buf + strlen (buf);
			int		remain_base = sizeof(buf) - strlen(buf);

			value = SPI_fname(tupdesc, i);
			type = SPI_gettype(tupdesc, i);
			snprintf(cur_loc, remain_base, " %4s(%s)%s", value, type, inter);
		}
		output_str("%5c %s", ' ', buf);
		output_str("%5c --------------------------------------------------", ' ');

        for (j = 0; j < rows; j++)
        {
            HeapTuple	tuple = tuptable->vals[j];

			memset(buf, 0, 8192);
            for (i = 1; i <= tupdesc->natts; i++)
			{
				char	*value = NULL;
				char	*inter = (i == tupdesc->natts) ? " " : " |";
				char	*cur_loc = buf + strlen (buf);
				int		remain_base = sizeof(buf) - strlen(buf);

				value = SPI_getvalue(tuple, tupdesc, i);
                snprintf(cur_loc, remain_base, " %10s%s", value, inter);
			}
			output_str("%5c %s", ' ', buf);
        }
    }
}


/*
 * A function package for SPI_exec, for the purpose we can
 * out put sql, and output exec result.
 */
static void
my_exec_spi(char *sql)
{
	int		ret = 0;

	output_str("SQL: %s", sql);
	ret = SPI_exec(sql, 0);
	outExecResult(ret);
}

static void
my_exec_spi_with_try(char *sql)
{
	int					ret = 0;
	MemoryContext		oldcontext = NULL;
	ResourceOwner		oldowner = NULL;

	oldcontext = CurrentMemoryContext;
	oldowner = CurrentResourceOwner;

	tx_begin(oldcontext);
	PG_TRY();
	{
		output_str("SQL: %s", sql);
		ret = SPI_exec(sql, 0);
		outExecResult(ret);
		tx_commit(oldcontext, oldowner);
	}
	PG_CATCH();
	{
		output_str("WARNING:Fail to exec %s, will ignore it.", sql);
		tx_abort(oldcontext, oldowner);
	}
	PG_END_TRY();
}

//////////////////////////////////////////////////////////////
// Some function for transaction test of SPI module
//////////////////////////////////////////////////////////////

static void
tx_begin(MemoryContext	oldcontext)
{
	output_str("SQL: BEGIN");
	BeginInternalSubTransaction(NULL);
	MemoryContextSwitchTo(oldcontext);
}

static void
tx_commit(MemoryContext	oldcontext, ResourceOwner oldowner)
{
	output_str("SQL: COMMIT");
	ReleaseCurrentSubTransaction();
	MemoryContextSwitchTo(oldcontext);
	CurrentResourceOwner = oldowner;
}

static void
tx_abort(MemoryContext	oldcontext, ResourceOwner oldowner)
{
	output_str("SQL: ABORT");
	RollbackAndReleaseCurrentSubTransaction();
	MemoryContextSwitchTo(oldcontext);
	CurrentResourceOwner = oldowner;
}

static void
spi_transaction_test(void)
{
	char	*sql10 = "INSERT INTO t1 VALUES(1,1,'I will be commited')";
	char	*sql11 = "INSERT INTO t1 VALUES(2,2,'I will be abort')";
	MemoryContext		oldcontext = NULL;
	ResourceOwner		oldowner = NULL;

	oldcontext = CurrentMemoryContext;
	oldowner = CurrentResourceOwner;
	
	SPI_connect();

	tx_begin(oldcontext);
	my_exec_spi(sql10);
	my_exec_spi(sql2);
	tx_commit(oldcontext, oldowner);

	tx_begin(oldcontext);
	my_exec_spi(sql11);
	my_exec_spi(sql2);
	tx_abort(oldcontext, oldowner);

	my_exec_spi(sql2);

	SPI_finish();
}

////////////////TRANSACTION TEST CODE END///////////////////////////////

static void
spi_with_args_test(void)
{
	char	*sql10 = "INSERT INTO t1 VALUES($1,$2,$3)";
	int		nargs = 3;
	Oid		argtypes[3] = {23, 23, 25};
	char 	nulls[3] = {' ', ' ', ' '};
	Datum	values[3];
	char	*text_value = "I am inserted by SPI_execute_with_args";

	values[0] = Int32GetDatum(100);
	values[1] = Int32GetDatum(100);
	values[2] = PointerGetDatum(cstring_to_text(text_value));

	SPI_connect();

	output_str("SQL:%s", sql10);
	SPI_execute_with_args(sql10, nargs, argtypes, values, nulls, false, 0);
	my_exec_spi(sql2);

	SPI_finish();
}

static void
spi_prepare_test(void)
{
	SPIPlanPtr	SPIpp = NULL;
	char	*sql10 = "INSERT INTO t1 VALUES(101, 101 , 'I am inserted by SPI prepare')";

	SPI_connect();

	output_str("SQL:%s", sql10);
	SPIpp = SPI_prepare(sql10, 0, NULL);
	SPI_execute_plan(SPIpp, NULL, NULL, false, 0);

	my_exec_spi(sql2);
	SPI_finish();
}

static void
spi_cursor_test(void)
{
	char		*sql10 = "INSERT INTO t1 SELECT generate_series(1,20),0,'Insert by generate_series'";
	char		*sql11 = "SELECT * FROM t1";
	Portal 		pt = NULL;
	SPIPlanPtr	SPIpp = NULL;
	
	SPI_connect();

	my_exec_spi(sql10);

	SPIpp = SPI_prepare(sql11, 0, NULL);
	output_str("SQL:Create Cursor");
	pt = SPI_cursor_open("Test_pt", SPIpp, NULL, NULL, false);

	output_str("SQL:Cursor Fetch3");
	SPI_cursor_fetch(pt, true, 3);
	outExecResult(1);
	output_str("SQL:Cursor Fetch3");
	SPI_cursor_fetch(pt, true, 3);
	outExecResult(1);

	SPI_cursor_close(pt);
	SPI_finish();
}



static void
spi_error_handle(void)
{
	char	*sql10 = "INSERT INTO t1 VALUES(1,1,'Test error handle')";
	char	*sql11 = "INSERT INTO t2 VALUES(1,1,'Where is the table t2?')";
	char	*sql12 = "INSERT INTO t1 VALUES(2,2,'Test error handle')";

	SPI_connect();

	my_exec_spi_with_try(sql10);
	my_exec_spi_with_try(sql11);
	my_exec_spi_with_try(sql12);

	my_exec_spi(sql2);

	SPI_finish();

}

static void
spi_simple_test(void)
{	
	char	*sql10 = "INSERT INTO t1 VALUES(1,1,'I will be deleted')";
	char	*sql11 = "INSERT INTO t1 VALUES(2,2,'I will be updated')";
	char	*sql12 = "INSERT INTO t1 VALUES(3,2,'I will be there all the time')";
	char	*sql13 = "UPDATE t1 SET k = 'updated from i=2' where i = 2";
	char	*sql14 = "DELETE FROM t1 WHERE i = 1";
	
	SPI_connect();

	my_exec_spi(sql10);
	my_exec_spi(sql11);
	my_exec_spi(sql12);
	my_exec_spi(sql13);
	my_exec_spi(sql14);
	my_exec_spi(sql2);

	SPI_finish();
}

Datum
test_spt_feature(PG_FUNCTION_ARGS)
{
	int	type = PG_GETARG_INT32(0);

	SPI_connect();
	my_exec_spi(sql0);
	my_exec_spi(sql1);
	SPI_finish();

	if(0 == type)
		spi_simple_test();
	else if(1 == type)
		spi_transaction_test();
	else if(2 == type)
		spi_with_args_test();
	else if(3 == type)
		spi_prepare_test();
	else if(4 == type)
		spi_cursor_test();
	else if(5 == type)
		spi_error_handle();

	PG_RETURN_CSTRING("SPI TEST!");
}

Datum
autoinc(PG_FUNCTION_ARGS)
{
	TriggerData *trigdata = (TriggerData *) fcinfo->context;
	Trigger    *trigger;		/* to get trigger name */
	int			nargs;			/* # of arguments */
	int		   *chattrs;		/* attnums of attributes to change */
	int			chnattrs = 0;	/* # of above */
	Datum	   *newvals;		/* vals of above */
	bool	   *newnulls;		/* null flags for above */
	char	  **args;			/* arguments */
	char	   *relname;		/* triggered relation name */
	Relation	rel;			/* triggered relation */
	HeapTuple	rettuple = NULL;
	TupleDesc	tupdesc;		/* tuple description */
	bool		isnull;
	int			i;

	if (!CALLED_AS_TRIGGER(fcinfo))
		/* internal error */
		elog(ERROR, "not fired by trigger manager");
	if (!TRIGGER_FIRED_FOR_ROW(trigdata->tg_event))
		/* internal error */
		elog(ERROR, "must be fired for row");
	if (!TRIGGER_FIRED_BEFORE(trigdata->tg_event))
		/* internal error */
		elog(ERROR, "must be fired before event");

	if (TRIGGER_FIRED_BY_INSERT(trigdata->tg_event))
		rettuple = trigdata->tg_trigtuple;
	else if (TRIGGER_FIRED_BY_UPDATE(trigdata->tg_event))
		rettuple = trigdata->tg_newtuple;
	else
		/* internal error */
		elog(ERROR, "cannot process DELETE events");

	rel = trigdata->tg_relation;
	relname = SPI_getrelname(rel);

	trigger = trigdata->tg_trigger;

	nargs = trigger->tgnargs;
	if (nargs <= 0 || nargs % 2 != 0)
		/* internal error */
		elog(ERROR, "autoinc (%s): even number gt 0 of arguments was expected", relname);

	args = trigger->tgargs;
	tupdesc = rel->rd_att;

	chattrs = (int *) palloc(nargs / 2 * sizeof(int));
	newvals = (Datum *) palloc(nargs / 2 * sizeof(Datum));
	newnulls = (bool *) palloc(nargs / 2 * sizeof(bool));

	for (i = 0; i < nargs;)
	{
		int			attnum = SPI_fnumber(tupdesc, args[i]);
		int32		val;
		Datum		seqname;

		if (attnum <= 0)
			ereport(ERROR,
					(errcode(ERRCODE_TRIGGERED_ACTION_EXCEPTION),
					 errmsg("\"%s\" has no attribute \"%s\"",
							relname, args[i])));

		if (SPI_gettypeid(tupdesc, attnum) != INT4OID)
			ereport(ERROR,
					(errcode(ERRCODE_TRIGGERED_ACTION_EXCEPTION),
					 errmsg("attribute \"%s\" of \"%s\" must be type INT4",
							args[i], relname)));

		val = DatumGetInt32(SPI_getbinval(rettuple, tupdesc, attnum, &isnull));

		if (!isnull && val != 0)
		{
			i += 2;
			continue;
		}

		i++;
		chattrs[chnattrs] = attnum;
		seqname = CStringGetTextDatum(args[i]);
		newvals[chnattrs] = DirectFunctionCall1(nextval, seqname);
		/* nextval now returns int64; coerce down to int32 */
		newvals[chnattrs] = Int32GetDatum((int32) DatumGetInt64(newvals[chnattrs]));
		if (DatumGetInt32(newvals[chnattrs]) == 0)
		{
			newvals[chnattrs] = DirectFunctionCall1(nextval, seqname);
			newvals[chnattrs] = Int32GetDatum((int32) DatumGetInt64(newvals[chnattrs]));
		}
		newnulls[chnattrs] = false;
		pfree(DatumGetTextPP(seqname));
		chnattrs++;
		i++;
	}

	if (chnattrs > 0)
	{
		rettuple = heap_modify_tuple_by_cols(rettuple, tupdesc,
											 chnattrs, chattrs,
											 newvals, newnulls);
	}

	pfree(relname);
	pfree(chattrs);
	pfree(newvals);
	pfree(newnulls);

	return PointerGetDatum(rettuple);
}
